#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
🎯 最终稳定版 YOLOv8n INT8 量化脚本（适配 nncf 新版 API + torch==2.5.0）
📌 特性：
   - 不依赖 TensorDataset
   - 使用 nncf.Dataset 正确封装
   - 兼容 ignored_scope（单数）和 TargetDevice 枚举
   - 可在 openvino-dev 非标准构建下运行
"""

import os
import sys
import cv2
import numpy as np
from pathlib import Path
from openvino import Core, save_model
import nncf
import logging
import random
import argparse

# ==================== 配置日志 ====================
def setup_logger():
    logger = logging.getLogger("INT8_Quantizer")
    logger.setLevel(logging.INFO)
    if not logger.handlers:
        handler = logging.StreamHandler()
        formatter = logging.Formatter(
            fmt='%(asctime)s [%(levelname)s] %(message)s',
            datefmt='%Y-%m-%d %H:%M:%S'
        )
        handler.setFormatter(formatter)
        logger.addHandler(handler)
    return logger

logger = setup_logger()


# ==================== 参数解析 ====================
def parse_args():
    parser = argparse.ArgumentParser(description="YOLOv8n OpenVINO 模型 INT8 量化工具 (新版 API)")

    parser.add_argument("--model", type=str, required=True,
                        help="FP32 OpenVINO 模型路径（.xml 文件）")
    parser.add_argument("--calib_dir", type=str, default="calibration_data",
                        help="校准图像目录路径（默认: calibration_data/）")
    parser.add_argument("--output", type=str, default="yolov8n_int8",
                        help="输出目录（默认: yolov8n_int8/）")
    parser.add_argument("--imgsz", type=int, default=480,
                        help="模型输入尺寸（默认: 480）")
    parser.add_argument("--samples", type=int, default=121,
                        help="用于校准的图像数量（默认: 121）")
    parser.add_argument("--seed", type=int, default=42,
                        help="随机种子（保证可复现）")

    return parser.parse_args()


# ==================== 主函数 ====================
def main():
    args = parse_args()

    # 设置随机种子
    random.seed(args.seed)
    np.random.seed(args.seed)

    logger.info("🚀 开始 INT8 量化流程 (适配新版 NNCF API)")
    logger.info(f"📁 模型文件: {args.model}")
    logger.info(f"🖼️  校准数据: {args.calib_dir} (使用最多 {args.samples} 张图)")
    logger.info(f"📐 输入尺寸: {args.imgsz}x{args.imgsz}")

    # -------------------------------
    # 1. 加载 FP32 模型
    # -------------------------------
    model_path = Path(args.model)
    if not model_path.exists():
        logger.error(f"❌ 模型文件不存在: {model_path}")
        sys.exit(1)

    try:
        core = Core()
        model = core.read_model(model_path)
        input_name = model.input().get_any_name()
        logger.info(f"✅ 成功加载 OpenVINO 模型，输入名: '{input_name}'")
    except Exception as e:
        logger.error(f"❌ 读取模型失败: {str(e)}")
        sys.exit(1)

    # -------------------------------
    # 2. 构建校准数据集
    # -------------------------------
    calib_dir = Path(args.calib_dir)
    if not calib_dir.is_dir():
        logger.error(f"❌ 校准目录不存在或不是文件夹: {calib_dir}")
        sys.exit(1)

    supported_exts = ["*.jpg", "*.jpeg", "*.png", "*.bmp"]
    all_images = []
    for ext in supported_exts:
        all_images.extend(calib_dir.glob(ext))

    if len(all_images) == 0:
        logger.error(f"❌ 在 {calib_dir} 中未找到图像文件")
        sys.exit(1)

    selected_images = random.sample(all_images, min(args.samples, len(all_images)))
    logger.info(f"📊 使用 {len(selected_images)} 张图像进行校准")

    # -------------------------------
    # 3. 创建 nncf.Dataset（关键：必须用 Dataset 包装）
    # -------------------------------
    def transform_fn(img_path):
        img = cv2.imread(str(img_path))
        if img is None:
            raise ValueError(f"无法读取图像: {img_path}")
        resized = cv2.resize(img, (args.imgsz, args.imgsz), interpolation=cv2.INTER_AREA)
        rgb = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)
        tensor = np.expand_dims(rgb.transpose(2, 0, 1), axis=0).astype(np.float32) / 255.0
        return {input_name: tensor}

    try:
        dataset = nncf.Dataset(selected_images, transform_fn)
        logger.info("✅ 校准数据集已通过 nncf.Dataset 封装")
    except Exception as e:
        logger.error(f"❌ 数据集封装失败: {str(e)}")
        sys.exit(1)

    # -------------------------------
    # 4. 执行 INT8 量化（使用新版 API）
    # -------------------------------
    logger.info("⚡ 开始量化（可能需要 2–5 分钟）...")

    try:
        from nncf import TargetDevice, QuantizationPreset
    except ImportError:
        logger.error("❌ 无法导入 TargetDevice 或 QuantizationPreset，请检查 nncf 是否为 2.7+")
        sys.exit(1)

    try:
        from nncf.scopes import IgnoredScope
    except ImportError:
        try:
            from nncf import IgnoredScope
        except ImportError:
            logger.error("❌ 无法导入 IgnoredScope，请更新 nncf")
            sys.exit(1)

    try:
        quantized_model = nncf.quantize(
            model,
            calibration_dataset=dataset,
            preset=QuantizationPreset.PERFORMANCE,
            target_device=TargetDevice.CPU,
            subset_size=len(selected_images),
            ignored_scope=IgnoredScope(names=["torchvision::nms/NonMaxSuppression"])
        )
        logger.info("✅ 量化成功")
    except TypeError as e:
        if "unexpected keyword argument 'ignored_scope'" in str(e):
            logger.error("❌ 当前版本不支持 ignored_scope，请尝试旧版 API")
            logger.info("💡 提示：你可能需要使用 openvino-dev==2024.0.0")
            sys.exit(1)
        else:
            raise
    except Exception as e:
        logger.error(f"❌ 量化失败: {str(e)}")
        sys.exit(1)

    # -------------------------------
    # 5. 保存 INT8 模型
    # -------------------------------
    output_dir = Path(args.output)
    output_dir.mkdir(parents=True, exist_ok=True)
    output_xml = output_dir / "yolov8n_int8.xml"

    try:
        save_model(quantized_model, str(output_xml))
        logger.info(f"🎉 INT8 模型已保存:")
        logger.info(f"   ├── {output_xml}")
        logger.info(f"   └── {output_xml.with_suffix('.bin')}")

        fp32_mb = model_path.stat().st_size / 1024 / 1024
        int8_mb = output_xml.stat().st_size / 1024 / 1024
        logger.info(f"📦 模型大小: FP32={fp32_mb:.1f}MB → INT8={int8_mb:.1f}MB")
        logger.info(f"🚀 预期推理速度提升: ~2–3x")

    except Exception as e:
        logger.error(f"❌ 保存模型失败: {str(e)}")
        sys.exit(1)

    logger.info("✨ 量化完成！现在可以用 INT8 模型替换原模型以显著提升推理性能！")


if __name__ == "__main__":
    main()